#include "antiShape.hpp"

AntiShapeParam::AntiShapeParam(){}
void AntiShapeParam::load(std::string filePath){
        Json::Reader reader;
        Json::Value root;
        std::ifstream infile(filePath);
        if(reader.parse(infile, root)){
            cropRatio = root["cropRatio"].asFloat();
            alpha = root["alpha"].asFloat();
            alphaMin = root["alphaMin"].asFloat();
            alphaMax = root["alphaMax"].asFloat();
            beta = root["beta"].asFloat();
            gamma = root["gamma"].asFloat();
            damping = root["damping"].asFloat();
            innerPaddingRatio = root["innerPaddingRatio"].asFloat();
            useSmooth = root["useSmooth"].asInt();
            lookaheadNo = root["lookaheadNo"].asInt();
            dx = root["dx"].asFloat();
            dy = root["dy"].asFloat();
            dealtWidth = root["dealtWidth"].asInt();
            dealtHeight = root["dealtHeight"].asInt();
            frameWidth = root["frameWidth"].asInt();
            frameHeight = root["frameHeight"].asInt();
            std::cout << "Loading config file successfully!" << std::endl;
        }
        else{
            std::cout << "Loading config file failed!" << std::endl;
        }
    }

    AntiShape::AntiShape(std::string configFile){
        this->config.load(configFile); 
        this->preP      = Eigen::MatrixXd::Identity(3,3);  
        this->preV      = Eigen::MatrixXd::Identity(3,3);  
        this->preDeltaP = Eigen::MatrixXd::Identity(3,3);   
        this->preDeltaV = Eigen::MatrixXd::Identity(3,3);  
        this->prePtV2P  = Eigen::MatrixXd::Identity(3,3);    

        float cropRatio = config.cropRatio;
        float innerPaddingRatio = config.innerPaddingRatio;
        float outerPaddingRatio = cropRatio * (1 - innerPaddingRatio);
        this->ScaleMat = Eigen::MatrixXd::Identity(3,3);  
        this->ScaleMat(1,1) = this->config.dy / this->config.dx;
        this->config.dealtHeight = (int)(this->config.dealtHeight * this->config.dy / this->config.dx);

        int outerWidth = (int)(cropRatio * this->config.dealtWidth);
        int outerHeight = (int)(cropRatio * this->config.dealtHeight);
        this->outerDist = outerHeight;

        //generate roi
        this->cropRoi = Eigen::MatrixXd::Zero(4,3); 
        float cropRoiX1 = cropRatio * (this->config.dealtWidth - 1.0) - (this->config.dealtWidth - 1.0) / 2.0;
        float cropRoiX2 = (1.0 - cropRatio) * (this->config.dealtWidth-1.0) - (this->config.dealtWidth - 1.0) / 2.0;
        float cropRoiY1 = cropRatio * (this->config.dealtHeight-1.0) - (this->config.dealtHeight - 1.0) / 2.0;
        float cropRoiY2 = (1.0 - cropRatio) * (this->config.dealtHeight-1.0) - (this->config.dealtHeight - 1.0) / 2.0;
        this->cropRoi << cropRoiX1, cropRoiY1, 1.0,
                         cropRoiX2, cropRoiY1, 1.0,
                         cropRoiX2, cropRoiY2, 1.0,
                         cropRoiX1, cropRoiY2, 1.0;

        this->frameCropRoiX1 = (int)(cropRatio * (this->config.frameWidth-1.0));
        this->frameCropRoiX2 = (int)((1.0 - cropRatio) * (this->config.frameWidth-1.0));
        this->frameCropRoiY1 = (int)(cropRatio * (this->config.frameHeight-1.0));
        this->frameCropRoiY2 = (int)((1.0 - cropRatio) * (this->config.frameHeight-1.0));

        //generate inner region
        this->innerRoi = Eigen::MatrixXd::Zero(4,3); 
        float innerRoiX1 = outerPaddingRatio * (this->config.dealtWidth-1.0) - (this->config.dealtWidth - 1.0) / 2.0;
        float innerRoiX2 = (1.0 - outerPaddingRatio) * (this->config.dealtWidth-1.0) - (this->config.dealtWidth - 1.0) / 2.0;
        float innerRoiY1 = outerPaddingRatio * (this->config.dealtHeight-1.0) - (this->config.dealtHeight - 1.0) / 2.0;
        float innerRoiY2 = (1.0 - outerPaddingRatio) * (this->config.dealtHeight-1.0) - (this->config.dealtHeight - 1.0) / 2.0;
        this->innerRoi << innerRoiX1, innerRoiY1, 1.0,
                          innerRoiX2, innerRoiY1, 1.0,
                          innerRoiX2, innerRoiY2, 1.0,
                          innerRoiX1, innerRoiY2, 1.0;

        // generate outer region
        this->outerRoi = Eigen::MatrixXd::Zero(4,3); 
        float outerRoiX1 = 0.0 - (this->config.dealtWidth - 1.0) / 2.0;
        float outerRoiX2 = this->config.dealtWidth - 1.0 - (this->config.dealtWidth - 1) / 2.0;
        float outerRoiY1 = 0.0 - (this->config.dealtHeight - 1.0) / 2.0;
        float outerRoiY2 = this->config.dealtHeight - 1.0 - (this->config.dealtHeight - 1) / 2.0;
        this->outerRoi << outerRoiX1, outerRoiY1, 1.0,
                          outerRoiX2, outerRoiY1, 1.0,
                          outerRoiX2, outerRoiY2, 1.0,
                          outerRoiX1, outerRoiY2, 1.0;                
    }    


    void AntiShape::put(cv::Mat& img){
        if(queue.size() < 500){
            queue.push(img);
        }
    }

    void AntiShape::get_xya(const Eigen::MatrixXd& rotM, tranParam& param){
        param.x = rotM(0, 2);
        param.y = rotM(1, 2);
        param.a = std::atan((float)(rotM(1, 0)/rotM(0, 0)));
    }

    void AntiShape::slerp_like(Eigen::MatrixXd& rotM, const Eigen::MatrixXd& rotM0, 
                    const Eigen::MatrixXd& rotM1, float alpha){
        tranParam par0, par1, par;
        this->get_xya(rotM0, par0);
        this->get_xya(rotM1, par1);
        par.x = (1.0 - alpha) * par0.x + alpha * par1.x;
        par.y = (1.0 - alpha) * par0.y + alpha * par1.y;
        par.a = (1.0 - alpha) * par0.a + alpha * par1.a;
        float sinA = std::sin(par.a);
        float cosA = std::cos(par.a);
        rotM << cosA, -sinA, par.x,
                sinA,  cosA, par.y,
                0,     0,    1; 
    }
    void AntiShape::cal_v_cam_velocity(Eigen::MatrixXd& currDeltaV, const Eigen::MatrixXd& preDeltaP, 
                                                         const Eigen::MatrixXd& curPtV2P){
        Eigen::MatrixXd preDeltaPInV = (curPtV2P.inverse() * preDeltaP) * this->prePtV2P;
        this->slerp_like(currDeltaV, preDeltaP, this->preDeltaV, this->config.alpha);
    }

    bool AntiShape::boundary_detection(int& uncompensatedDist, const Eigen::MatrixXd& curPtV2P){
        // Eigen::MatrixXd cropRoiV2P = curPtV2P * this->cropRoi.transpose();
        // cropRoiV2P = cropRoiV2P.transpose();
        Eigen::MatrixXd cropRoiV2P = this->cropRoi * curPtV2P.transpose();
        bool isInFrame = true;
        int dist;
        uncompensatedDist = 9999;
        for(int i = 0; i < cropRoiV2P.rows(); i++){
            if((abs(cropRoiV2P(i,0)) <= (this->config.dealtWidth -  1.0) / 2.0) &&
               (abs(cropRoiV2P(i,1)) <= (this->config.dealtHeight - 1.0) / 2.0)){
                isInFrame = isInFrame && true;
            }else{
                isInFrame = isInFrame && false;
            }
            dist = std::min((int)((this->config.dealtWidth - 1.0)/2.0 + cropRoiV2P(i,0)), 
                            (int)((this->config.dealtWidth - 1.0)/2.0 - cropRoiV2P(i,0)));
            if(dist < uncompensatedDist){
                uncompensatedDist = dist;
            }
            
            dist = std::min((int)((this->config.dealtHeight - 1.0)/2.0 + cropRoiV2P(i,1)), 
                            (int)((this->config.dealtHeight - 1.0)/2.0 - cropRoiV2P(i,1)));
            if(dist < uncompensatedDist){
                uncompensatedDist = dist;
            }
        }
        if(isInFrame){
            return false;
        }
        uncompensatedDist = 0;
        return true;
    }

    bool AntiShape::update_filter(Eigen::MatrixXd& currDeltaV,bool isOutOfBoundary, const Eigen::MatrixXd& curPtV2P, int uncompensatedDist){
        currDeltaV = Eigen::MatrixXd::Identity(3,3);
        bool isInInnerRegion = true;
        Eigen::MatrixXd cropRoiV2P = this->cropRoi * curPtV2P.transpose();
        for(int i = 0; i < cropRoiV2P.rows(); i++){
            if ((cropRoiV2P(i,0) >= this->innerRoi(0,0)) && (cropRoiV2P(i,0) <= this->innerRoi(2,0)) &&
                (cropRoiV2P(i,1) >= this->innerRoi(0,1)) && (cropRoiV2P(i,1) <= this->innerRoi(2,1))){
                isInInnerRegion = isInInnerRegion && true;
            }else{
                isInInnerRegion = isInInnerRegion && false;
            }
        }
        if(isInInnerRegion){
            Eigen::MatrixXd rotM = Eigen::MatrixXd::Identity(3,3);
            this->slerp_like(currDeltaV, rotM, this->preDeltaV, this->config.damping);
        }else{
            float alphaRatio = std::pow((float)uncompensatedDist/(float)(this->outerDist), this->config.beta);
            alphaRatio = std::fmin(std::fmax(alphaRatio, 0.0) , 1.0);
            float alpha = alphaRatio * this->config.alphaMax + (1 - alphaRatio) * this->config.alphaMin;
            this->config.alpha  = 0.95 * this->config.alpha + 0.05*alpha;
            this->cal_v_cam_velocity(currDeltaV, this->preDeltaP, curPtV2P);  
        }
        return isInInnerRegion;  
    }

    void AntiShape::updateRotationMatrix(Eigen::MatrixXd& rotM){
        double a = std::atan((double)(rotM(1, 0)/rotM(0, 0)));
        double cosA = std::cos(a);
        double sinA = std::sin(a);
        rotM(0,0) = cosA;
        rotM(1,1) = cosA;
        rotM(1,0) = sinA;     
        rotM(0,1) = -sinA;   
    }

    
    void AntiShape::smooth(Eigen::MatrixXd& smoothedRotM, const std::deque<Eigen::MatrixXd>& trajMatrixs){ 
        tranParam par, parSum;
        parSum.x = 0.0;
        parSum.y = 0.0;
        parSum.a = 0.0;
        for(int i = 0; i < trajMatrixs.size(); i++){
            get_xya(trajMatrixs[i], par);
            parSum.x += par.x;
            parSum.y += par.y;
            parSum.a += par.a;
        }
        parSum.x /= trajMatrixs.size(); 
        parSum.y /= trajMatrixs.size(); 
        parSum.a /= trajMatrixs.size();
        float sinA = std::sin(parSum.a);
        float cosA = std::cos(parSum.a);
        smoothedRotM << cosA, -sinA, parSum.x,
                        sinA,  cosA, parSum.y,
                        0,     0,    1;
    }

    void AntiShape::getMovement(Eigen::MatrixXd& rotM, const cv::Mat& imgCurrGray, const cv::Mat& imgPrevGray){
        // Detect features in previous frame
 
        std::vector<cv::Point2f> have_corners, prevPts, currPts;
        // efficientGoodFeaturesToTrack(imgPrevGray, have_corners, prevPts, 200, 30);
        cv::goodFeaturesToTrack(imgPrevGray, prevPts, 200, 0.01, 30);//, blockSize=3

        // Calculate optical flow (i.e. track feature points)
        if(prevPts.size() < 10){
            rotM = Eigen::MatrixXd::Identity(3,3);
            return;
        }

        std::vector<uchar> status;
        std::vector<float> err;
        cv::calcOpticalFlowPyrLK(imgPrevGray, imgCurrGray, prevPts, currPts,status, err);
    
        // # Sanity check 
        assert(prevPts.size() == currPts.size());

        int h = imgCurrGray.rows;
        int w = imgCurrGray.cols;

        for(int i = status.size()-1; i >= 0; i--){
            if(status[i] == 1){
                prevPts[i].x -= (w-1.0)/2.0;
                prevPts[i].y -= (h-1.0)/2.0;
                currPts[i].x -= (w-1.0)/2.0;
                currPts[i].y -= (h-1.0)/2.0;
            }else{
            // Filter only valid points
                prevPts.erase(prevPts.begin()+i);
                currPts.erase(currPts.begin()+i);
            }
        } 

        if(prevPts.size() < 10){
            rotM = Eigen::MatrixXd::Identity(3,3);
            return;
        }

        // # Find transformation matrix
        cv::Mat m = cv::estimateAffinePartial2D(prevPts, currPts);
        cv::cv2eigen(m,rotM);
        rotM.conservativeResize(rotM.rows()+1, rotM.cols());
        rotM.row(rotM.rows() - 1) << 0, 0, 1;
    }

    void AntiShape::get(cv::Mat& frameAffined, cv::Mat& frame){
        // start = time.time() ！！！！！！！！！！！！！！
        if(this->queue.size() <= 0) return;
        cv::Mat image = this->queue.front();
        this->queue.pop();
        if (this->imgPrev.empty()){
            cv::resize(image, this->imgPrev, cv::Size(this->config.dealtWidth,this->config.dealtHeight));
            if(this->config.lookaheadNo > 0){  
                //push to buffer   
                this->PBuffer.push_back(this->preP);
                this->VBuffer.push_back(this->preV);
                this->deltaVBuffer.push_back(this->preDeltaV);
                this->deltaPBuffer.push_back(Eigen::MatrixXd::Identity(3,3));
                this->imageBuffer.push_back(image);
            }
            return;
        }
        cv::Mat resizedImg, imgCurrGray, imgPrevGray;
        cv::resize(image, resizedImg, cv::Size(this->config.dealtWidth,this->config.dealtHeight));
        cv::cvtColor(resizedImg, imgCurrGray, cv::COLOR_BGR2GRAY);
        cv::cvtColor(this->imgPrev, imgPrevGray, cv::COLOR_BGR2GRAY);

        // calculate physical camera velocity (previous) and trajectory (current)
        this->getMovement(this->preDeltaP, imgCurrGray, imgPrevGray);
        Eigen::MatrixXd curP = this->preDeltaP * this->preP;
        this->updateRotationMatrix(curP);

        // calculate virtual camera trajectory (current)
        Eigen::MatrixXd curV = this->preDeltaV * this->preV;
        this->updateRotationMatrix(curV);

        // calculate current transformation metrix (points from v_cam to p_cam)
        Eigen::MatrixXd curPtV2P;
        int outIdx = (int)((this->config.lookaheadNo - 1) / 2);
        if(this->config.lookaheadNo > 0){
            if(!this->config.useSmooth){
                if(this->deltaVBuffer.size() >= this->config.lookaheadNo){
                    Eigen::MatrixXd fowardDeltaV = this->deltaVBuffer[0];
                    for(int i = this->config.lookaheadNo - 1; i >= 0; i--){ 
                        this->slerp_like(fowardDeltaV, this->deltaPBuffer[i], fowardDeltaV, this->config.gamma); //!!!!!!!!!
                    }
                    this->VBuffer[1] = fowardDeltaV * this->VBuffer[0];
                    this->updateRotationMatrix(this->VBuffer[1]);
                    curPtV2P = this->PBuffer[1] * this->VBuffer[1].inverse();
                    this->updateRotationMatrix(curPtV2P);
                }else{
                    curPtV2P = Eigen::MatrixXd::Identity(3,3);
                }
            }else{
                if(this->PBuffer.size() >= this->config.lookaheadNo){
                    this->smooth(this->VBuffer[outIdx], this->PBuffer);
                    curPtV2P = this->PBuffer[outIdx] * this->VBuffer[outIdx].inverse();
                    this->updateRotationMatrix(curPtV2P);
                }else{
                    curPtV2P = Eigen::MatrixXd::Identity(3,3);
                }
            }
        }else{
            curPtV2P = curP * curV.inverse();
            this->updateRotationMatrix(curPtV2P);
        }
        // check if crop roi is out of frame
        int uncompensatedDist;
        bool isOutOfBoundary = this->boundary_detection(uncompensatedDist, curPtV2P);

        // update current virtual camera trajectory and current transformation metrix
        if(isOutOfBoundary){ 
            curPtV2P = this->prePtV2P;
            if(this->config.lookaheadNo > 0){
                if(!this->config.useSmooth){
                    this->VBuffer[1] = curPtV2P.inverse() * this->PBuffer[1];
                    this->updateRotationMatrix(this->VBuffer[1]);
                }
                else{  
                    this->VBuffer[outIdx] = curPtV2P.inverse() * this->PBuffer[outIdx];
                    this->updateRotationMatrix(this->VBuffer[outIdx]);
                }
            }else{
                curV = curPtV2P.inverse() * curP;
                this->updateRotationMatrix(curV);
            }
        }

        // calculate current virtual camera velocity
        Eigen::MatrixXd currDeltaV;
        bool isInInnerRegion = this->update_filter(currDeltaV, isOutOfBoundary, curPtV2P, uncompensatedDist);
        if ((this->config.lookaheadNo == 0) || 
            ((!this->config.useSmooth) && this->config.lookaheadNo > 0 && this->imageBuffer.size() > 1) ||
            (this->config.useSmooth && this->config.lookaheadNo > 0 && this->imageBuffer.size() >= this->config.lookaheadNo)){
            if(this->config.lookaheadNo > 0){
                if(!this->config.useSmooth){
                    frame = this->imageBuffer[1];
                }else{
                    frame = this->imageBuffer[outIdx];
                }
            }else{
                frame = image;
            }
            Eigen::MatrixXd curPtP2V = curPtV2P.inverse();
            this->updateRotationMatrix(curPtP2V);
            int h = frame.rows;
            int w = frame.cols;
            float dx = curPtP2V(0, 2);
            float dy = curPtP2V(1, 2);

            float da = std::atan((float)(curPtP2V(1, 0)/curPtP2V(0, 0)));
            // Reconstruct transformation matrix accordingly to new values
            cv::Point2f center((w-1)/2.0, (h-1)/2.0);
            cv::Mat m = cv::getRotationMatrix2D(center, -180.0*da/M_PI, 1.0);
            Eigen::MatrixXd rotM;
            cv::cv2eigen(m,rotM);
            rotM(0,2) += dx*w/this->config.dealtWidth;
            rotM(1,2) += dy*h/this->config.dealtHeight;
            // print(self.ScaleMat)
            rotM.conservativeResize(rotM.rows()+1, rotM.cols());
            rotM.row(rotM.rows() - 1) << 0, 0, 1;   

            rotM = this->ScaleMat.inverse() * rotM;
            this->updateRotationMatrix(rotM);
            rotM.conservativeResize(rotM.rows()-1, rotM.cols());
            cv::eigen2cv(rotM,m);

            // Apply affine wrapping to the given frame
            cv::warpAffine(frame, frameAffined, m, cv::Size(w, h));
            
            frameAffined = frameAffined(cv::Range(this->frameCropRoiY1, this->frameCropRoiY2),
                                        cv::Range(this->frameCropRoiX1, this->frameCropRoiX2));
            cv::resize(frameAffined,frameAffined,cv::Size(w,h));
        }
        if(this->config.lookaheadNo > 0){    
            // push to buffer    
            this->PBuffer.push_back(curP);
            this->VBuffer.push_back(curV);
            this->deltaVBuffer.push_back(currDeltaV);
            this->deltaPBuffer.push_back(this->preDeltaP);
            this->imageBuffer.push_back(image);
            if(this->PBuffer.size() > this->config.lookaheadNo) this->PBuffer.pop_front();
            if(this->VBuffer.size() > this->config.lookaheadNo) this->VBuffer.pop_front();
            if(this->deltaVBuffer.size() > this->config.lookaheadNo) this->deltaVBuffer.pop_front();
            if(this->deltaPBuffer.size() > this->config.lookaheadNo) this->deltaPBuffer.pop_front();
            if(this->imageBuffer.size() > this->config.lookaheadNo) this->imageBuffer.pop_front();
        }

        this->preP = curP;
        this->preV = curV;
        this->prePtV2P = curPtV2P;
        this->preDeltaV = currDeltaV;
        this->imgPrev = resizedImg;

        // # print("get: ", time.time()-start)
    }




// void dealPro(){
//     std::string configFile = "/home/hxzh/01-my_work/01-tools/01-good/python/antiShape/config.json";
//     // cv::VideoCapture cap("/home/hxzh/01-my_work/01-tools/00-others/temp/192.168.6.135_01_20231026134737911.mp4");
//     cv::VideoCapture cap("/home/hxzh/01-my_work/01-tools/00-others/temp/192.168.8.63_01_20230905150805435.mp4");

//     // nFrames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))  
//     // w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
//     // h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
//     int w = 640;
//     int h = 360;
//     double fps = cap.get(cv::CAP_PROP_FPS);
//     int fourcc = cv::VideoWriter::fourcc('M','J','P','G');
//     cv::VideoWriter out2("stabilized_video.avi", fourcc, fps, cv::Size(w, h*2));

//     AntiShape antiShape(configFile);
//     int count = 0;
//     std::string imgpath = "/home/hxzh/01-my_work/01-tools/00-others/temp2";
//     cv::Mat img;
//     struct timeval startTime, stopTime;
//     long elapsed;
//     while(1){
//         bool success = cap.read(img);
//         if(!success){
//             break;
//         }

//         // gettimeofday(&startTime, NULL);
//         cv::resize(img,img,cv::Size(w,h));
//         antiShape.put(img);
//         cv::Mat frameSta, frame;
//         antiShape.get(frameSta, frame);
//         count += 1;
//         std::cout << "count: " << count << std::endl;
//         if(count >= 300)
//             break;
//         if(!frameSta.empty()){
//             cv::Mat temp;
//             cv::vconcat(frame, frameSta, temp);

//             // if not qData.full():
//             //      qData.put(temp)
//             out2.write(temp);
//             //     # cv2.imwrite(imgpath+"/{:06d}.jpg".format(count),temp)
//             // allTime = time.time() - start
//             // gettimeofday(&stopTime, NULL);
//             // elapsed = (stopTime.tv_sec-startTime.tv_sec) * 1000000 + (stopTime.tv_usec - startTime.tv_usec);
//             // std::cout << "0: " << elapsed << "us" << std::endl;
//             // gettimeofday(&startTime, NULL);
//             // if allTime < 0.04: 
//             //     time.sleep(0.04 - allTime)
//         }
//     }
//     cap.release();
//     out2.release();
// }

// int main(void){



//     AntiShape test("/home/hxzh/01-my_work/01-tools/01-good/python/antiShape/config.json");
//     cv::Mat img0 = cv::imread("/home/hxzh/01-my_work/01-tools/00-others/temp2/000001.jpg",cv::IMREAD_COLOR);
//     cv::Mat img1 = cv::imread("/home/hxzh/01-my_work/01-tools/00-others/temp2/000002.jpg",cv::IMREAD_COLOR);


//     std::thread thread0(dealPro);
//     thread0.join();

//     cv::Mat img0Gray, img1Gray;
//     cv::cvtColor(img0,img0Gray,cv::COLOR_BGR2GRAY);
//     cv::cvtColor(img1,img1Gray,cv::COLOR_BGR2GRAY);
//     Eigen::MatrixXd rotM = Eigen::MatrixXd::Identity(3,3);  ;
//     // test.getMovement(rotM, img0Gray,img1Gray);

//     return 0;
// }
